#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
mosaic_tifs.py

将多个 GeoTIFF 进行拼接合并，输出一个新的 GeoTIFF。

特性：
- 支持输入目录或文件列表
- 自动对齐到统一 CRS（默认使用第一个影像的 CRS，或通过 --dst-crs 指定）
- 自动推断输出范围
- 支持设置 nodata 与 resampling 方法
- 保留压缩（LZW）

依赖：rasterio, numpy

用法示例：
1) 使用输入目录（递归搜索 .tif/.tiff）：
   python mosaic_tifs.py --input-dir ./tiles --output ./mosaic.tif

2) 指定文件列表：
   python mosaic_tifs.py --inputs a.tif b.tif c.tif --output ./mosaic.tif

3) 指定目标 CRS 与 nodata：
   python mosaic_tifs.py --input-dir ./tiles --output ./mosaic.tif --dst-crs EPSG:4326 --nodata 0
"""

import argparse
import os
from typing import List, Optional

import numpy as np
import rasterio
from rasterio.enums import Resampling
from rasterio.merge import merge
from rasterio.vrt import WarpedVRT


def select_valid_pixels(merged_data, new_data, merged_mask, new_mask, **kwargs):
    """
    自定义像素选择函数，用于处理重叠区域。
    如果两个数据源在同一位置都有有效像素，则选择new_data的像素值。
    如果只有merged_mask为True，保留merged_data的值。
    如果只有new_mask为True，使用new_data的值。
    """
    # 创建一个布尔掩码，表示merged_mask和new_mask都为True的位置
    valid_mask = np.logical_and(merged_mask, new_mask)

    # 创建一个临时掩码，用于标记merged_data中需要被更新的像素位置
    update_mask = np.copy(valid_mask)

    # 如果merged_mask为True但new_mask为False，保留merged_data的值
    np.logical_and(merged_mask, np.logical_not(new_mask), out=valid_mask)
    np.logical_or(update_mask, valid_mask, out=update_mask)

    # 如果new_mask为True但merged_mask为False，更新merged_data的值为new_data
    np.logical_and(np.logical_not(merged_mask), new_mask, out=valid_mask)
    np.logical_or(update_mask, valid_mask, out=update_mask)

    # 更新merged_data的值
    np.copyto(merged_data, new_data, where=update_mask, casting="unsafe")


def list_tifs(input_dir: str) -> List[str]:
	paths: List[str] = []
	for root, _dirs, files in os.walk(input_dir):
		for name in files:
			low = name.lower()
			if low.endswith('.tif') or low.endswith('.tiff'):
				paths.append(os.path.join(root, name))
	return paths


def build_src_list(paths: List[str]) -> List[rasterio.io.DatasetReader]:
	srcs = []
	for p in paths:
		srcs.append(rasterio.open(p))
	return srcs


def write_geotiff(path: str, array: np.ndarray, profile: dict) -> None:
	prof = profile.copy()
	prof.update({
		'count': array.shape[0],
		'dtype': np.dtype(array.dtype).name,
		'compress': 'lzw',
	})
	with rasterio.open(path, 'w', **prof) as dst:
		for i in range(array.shape[0]):
			dst.write(array[i], i + 1)


def mosaic(paths: List[str], output: str, dst_crs: Optional[str], nodata: Optional[float], resampling: str, merge_method: str, use_custom_merge: bool = False) -> str:
	if not paths:
		raise SystemExit('未提供任何输入影像')
	# 重投影重采样（用于 WarpedVRT）
	resampling_map = {
		'nearest': Resampling.nearest,
		'bilinear': Resampling.bilinear,
		'cubic': Resampling.cubic,
		'cubic_spline': Resampling.cubic_spline,
		'lanczos': Resampling.lanczos,
		'average': Resampling.average,
		'mode': Resampling.mode,
	}
	warp_resampling = resampling_map.get(resampling, Resampling.nearest)

	# 打开所有源
	srcs = build_src_list(paths)
	try:
		# 若指定目标 CRS，则用 WarpedVRT 虚拟重投影后再合并
		if dst_crs:
			vrts = []
			for s in srcs:
				vrt = WarpedVRT(
					s,
					crs=dst_crs,
					resampling=warp_resampling,
					nodata=nodata if nodata is not None else s.nodata,
				)
				vrts.append(vrt)
			inputs = vrts
		else:
			inputs = srcs

		# 合并
		if use_custom_merge:
			# 使用自定义的select_valid_pixels方法
			mosaic_arr, mosaic_transform = merge(
				inputs,
				nodata=nodata,
				method=select_valid_pixels,
				bounds=None,
			)
		else:
			# 使用默认的合并方法
			mosaic_arr, mosaic_transform = merge(
				inputs,
				nodata=nodata,
				method=merge_method,
				bounds=None,
			)
		# 基于首个输入构建 profile（若使用 VRT，其 CRS 已为目标 CRS）
		profile = inputs[0].profile if not isinstance(inputs[0], WarpedVRT) else srcs[0].profile
		# 写出时覆盖关键元数据
		profile.update({
			'transform': mosaic_transform,
			'height': mosaic_arr.shape[1],
			'width': mosaic_arr.shape[2],
			'crs': dst_crs or srcs[0].crs,
			'nodata': nodata if nodata is not None else srcs[0].nodata,
		})
		write_geotiff(output, mosaic_arr, profile)
		print(f'输出: {output}')
		return output
	finally:
		# 关闭数据集与 VRT
		if 'inputs' in locals():
			for ds in inputs:
				try:
					ds.close()
				except Exception:
					pass
		for s in srcs:
			try:
				s.close()
			except Exception:
				pass


def parse_args() -> argparse.Namespace:
	p = argparse.ArgumentParser(description='拼接多个 GeoTIFF 为一个影像')
	g = p.add_mutually_exclusive_group(required=True)
	g.add_argument('--input-dir', help='输入目录（递归搜索 .tif/.tiff）')
	g.add_argument('--inputs', nargs='+', help='输入文件列表')
	p.add_argument('--output', required=True, help='输出 GeoTIFF 路径')
	p.add_argument('--dst-crs', help='目标 CRS（如 EPSG:4326），默认使用第一个影像 CRS')
	p.add_argument('--nodata', type=float, help='输出 nodata 值；缺省使用源影像 nodata')
	p.add_argument('--resampling', default='nearest',
				choices=['nearest','bilinear','cubic','cubic_spline','lanczos','average','mode'],
				help='重采样方法，默认 nearest')
	p.add_argument('--merge-method', default='first',
				choices=['first','last','min','max','sum','count'],
				help='合并重叠像元的方法（传给 merge），默认 first')
	p.add_argument('--use-custom-merge', action='store_true',
				help='使用自定义的select_valid_pixels方法处理重叠区域')
	return p.parse_args()


def merge_tifs_in_directory(province_dir: str, output_file: str, use_custom_merge: bool = True) -> None:
    """
    合并指定目录中的所有 TIFF 文件，并将结果写入到输出文件中。
    这是基于您提供的示例代码的便捷函数。

    参数:
    province_dir (str): 包含 TIFF 文件的目录路径。
    output_file (str): 合并结果的输出文件路径。
    use_custom_merge (bool): 是否使用自定义的select_valid_pixels方法，默认为True。
    """
    # 获取目录中的所有TIFF文件
    paths = list_tifs(province_dir)
    
    if not paths:
        print("没有找到有效的 TIFF 文件")
        return

    # 使用mosaic函数进行合并
    try:
        mosaic(
            paths=paths, 
            output=output_file, 
            dst_crs='EPSG:4326',  # 默认使用WGS84坐标系
            nodata=0,  # 默认nodata值为0
            resampling='nearest',  # 使用最近邻重采样
            merge_method='first',  # 当不使用自定义合并时的默认方法
            use_custom_merge=use_custom_merge
        )
        print(f"合并完成: {output_file}")
    except Exception as e:
        print(f"合并过程中出现错误: {e}")


def main() -> None:
	args = parse_args()
	if args.input_dir:
		paths = list_tifs(args.input_dir)
	else:
		paths = args.inputs or []
	if not paths:
		raise SystemExit('未在输入中找到任何 .tif/.tiff 文件')
	mosaic(paths=paths, output=args.output, dst_crs=args.dst_crs, nodata=args.nodata, resampling=args.resampling, merge_method=args.merge_method, use_custom_merge=args.use_custom_merge)


if __name__ == '__main__':
	# 如果直接运行此脚本，可以使用便捷函数
	# 取消注释下面的代码来使用示例：
	# province_dir = r'G:\download\GLC_FCS30_2020'
	# output_file = r'H:\land30m2020\merged_output\GLC_FCS30_2020_noreplace.tif'
	# merge_tifs_in_directory(province_dir, output_file, use_custom_merge=True)
	
	# 否则使用命令行参数
	main()
